# -*- coding: utf-8 -*-
# @Author: lidongdong
# @time  : 18-12-16 上午10:15
# @file  : image2h5.py

import h5py
import cv2
import numpy as np
import glob
import re

from utils.skipthoughts import *


jpg_path = "/dl_data/flower-gan/Data/jpg"
caption_path = "/dl_data/flower-gan/Data/text_c10"


def load_image(image_filename):
    origin_img = cv2.imread(image_filename)
    width, height, _ = origin_img.shape
    if width > height:
        pad = (width - height) / 2
        croped_image = origin_img[pad: pad + height, :, :]
    elif width < height:
        pad = (height - width) / 2
        croped_image = origin_img[:, pad: pad + width, :]
    else:
        croped_image = origin_img

    scaled_image = cv2.resize(croped_image, dsize=(64, 64))
    return scaled_image


def load_caption(caption_filename):
    with open(caption_filename) as f:
        content = f.read()
        captions = content.strip().split("\n")
        assert len(captions) >= 5
        return captions[:5]


def main():
    image_filenames = os.listdir(jpg_path)
    image_filenames = sorted(image_filenames)

    caption_filenames = glob.glob(caption_path + "/class_*/image_*.txt")
    caption_filenames.sort(key=lambda x: int(re.match(".*image_(.*).txt", x).group(1)))

    h5file = h5py.File("/dl_data/flower-gan/Data/images.h5", "w")
    X = h5file.create_dataset("image", shape=(0, 64, 64, 3), maxshape=(None, 64, 64, 3), compression="gzip", chunks=(1, 64, 64, 3), dtype=np.float32)
    Y = h5file.create_dataset("caption", shape=(0, 5, 4800), maxshape=(None, 5, 4800), compression="gzip", chunks=(1, 5, 4800), dtype=np.float32)

    model = load_model()
    encoder = Encoder(model)

    for idx, (img, cap) in enumerate(zip(image_filenames, caption_filenames)):
        print "h5 {}".format(img)
        img = os.path.join(jpg_path, img)
        image_data = np.expand_dims(load_image(img), 0)
        X.resize(X.shape[0] + 1, axis=0)
        X[idx: idx+1, :, :, :] = (image_data - 128.0) / 128.0

        data = load_caption(cap)
        Y.resize(Y.shape[0] + 1, axis=0)
        captions = np.asarray(encoder.encode(data))
        Y[idx: idx+1, :, :] = captions

    h5file.close()


if __name__ == '__main__':
    main()
